import os
from pinecone import Pinecone, ServerlessSpec
from langchain_pinecone import PineconeVectorStore
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document

# 设置你的 API Key
PINECONE_API_KEY = "pcsk_49Aa66_HvcsTrsz9gTPzmog8ex9vumxtF4J4EWbjq22YsveZBx5XuDgCFoL7VZxjZRdoBp"


def pinecone_demo():
    print("🚀 开始 Pinecone 演示...")

    # 1. 初始化 Pinecone（新版正确方式）
    pc = Pinecone(api_key=PINECONE_API_KEY)

    # 2. 创建或连接索引
    index_name = "langchain-demo"

    # 检查索引是否存在
    existing_indexes = pc.list_indexes().names()
    print(f"现有索引: {existing_indexes}")

    if index_name not in existing_indexes:
        print(f"📦 创建新索引: {index_name}")
        pc.create_index(
            name=index_name,
            dimension=384,  # 必须与嵌入模型维度匹配
            metric="cosine",
            spec=ServerlessSpec(
                cloud="aws",
                region="us-east-1"
            )
        )
        print("⏳ 等待索引初始化...")
        # 等待索引就绪
        import time
        time.sleep(10)
    else:
        print(f"✅ 连接到现有索引: {index_name}")

    # 3. 准备示例数据
    documents = [
        Document(
            page_content="机器学习是人工智能的一个分支，使计算机能够从数据中学习",
            metadata={"category": "AI", "type": "definition"}
        ),
        Document(
            page_content="深度学习使用神经网络处理复杂模式识别任务",
            metadata={"category": "AI", "type": "technique"}
        ),
        Document(
            page_content="自然语言处理让计算机理解和生成人类语言",
            metadata={"category": "NLP", "type": "application"}
        ),
        Document(
            page_content="计算机视觉使机器能够识别和理解图像内容",
            metadata={"category": "CV", "type": "application"}
        ),
        Document(
            page_content="Pinecone 是高效的云端向量数据库服务",
            metadata={"category": "Tool", "type": "service"}
        )
    ]

    print(f"📚 准备 {len(documents)} 个文档")

    # 4. 创建嵌入模型
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    # 5. 上传文档到 Pinecone
    print("🔄 上传文档到 Pinecone...")
    vector_store = PineconeVectorStore.from_documents(
        documents=documents,
        embedding=embeddings,
        index_name=index_name
    )

    print("✅ 文档上传完成！")

    # 6. 测试搜索功能
    print("\n🔍 测试向量搜索:")

    test_queries = [
        "什么是机器学习？",
        "AI 有哪些应用领域？",
        "解释深度学习",
        "图像识别技术"
    ]

    for query in test_queries:
        print(f"\n❓ 查询: {query}")
        results = vector_store.similarity_search(query, k=2)

        for i, doc in enumerate(results):
            print(f"📄 结果 {i + 1}: {doc.page_content}")
            print(f"   标签: {doc.metadata}")
        print("-" * 60)


if __name__ == "__main__":
    pinecone_demo()